#include "ntfsutils.h"
#include <cstring>

namespace ntfs 
{

/******************************************************************************/
/*                                CNTFSVolume                                 */
/******************************************************************************/

CNTFSVolume::CNTFSVolume(const tstring fileName)
: CDisk(fileName)
{
    read(&mBootBlock, sizeof(BOOT_BLOCK));
      
    //Is volume NTFS
    UINT32 VolumeFormat = *(PUINT32)&(mBootBlock.Format);
    if (VolumeFormat != 0x5346544E)
        throw ;
    
    mBytesPerCluster = mBootBlock.BytesPerSector * mBootBlock.SectorsPerCluster;
    mMFTOffset = mBytesPerCluster * mBootBlock.MftStartLcn;
    mBytesPerFileRecord = mBootBlock.ClustersPerFileRecord < 0x80
    ? mBootBlock.ClustersPerFileRecord * mBootBlock.SectorsPerCluster * mBootBlock.BytesPerSector
    : 1 << (0x100 - mBootBlock.ClustersPerFileRecord);
}

CNTFSVolume::~CNTFSVolume()
{
}

SxSharedBuffer CNTFSVolume::readCluster(UINT64 LCN) const
{

    SxSharedBuffer memBuff(getClusterLength());
    read(memBuff, getClusterLength(), mBytesPerCluster*LCN);
    return memBuff;
}

/******************************************************************************/
/*                                CFileRecord                                 */
/******************************************************************************/

CFileRecord::CFileRecord(const SxSharedBuffer& MemBuffer):
    mBuffer(MemBuffer)
{
    LoadAttributes();
}

CFileRecord::CFileRecord(const CFileRecord& FileRecord):
    mBuffer(FileRecord.mBuffer)
{
    LoadAttributes();
}

CFileRecord& CFileRecord::operator =(const CFileRecord& FileRecord)
{
    if (&FileRecord != this)
    {
        mBuffer = FileRecord.mBuffer;
        LoadAttributes();
    }
    return *this;
}

void CFileRecord::LoadAttributes()
{
    PFILE_RECORD_HEADER fileRec = mBuffer;

    UINT32 offset = fileRec->AttributesOffset;
    for (int j=0; ; j++)
    {
        if (offset > fileRec->BytesAllocated)
            break;

        PATTRIBUTE pAttr = (PATTRIBUTE)IncPTR(fileRec, offset);
        if (pAttr->AttributeType == -1) break;

        CAttribute* pAttribute;
        if (pAttr->Nonresident)
            pAttribute = new CNonResidentAttribute(SxSharedBuffer(mBuffer, offset));
        else
            pAttribute = new CResidentAttribute(SxSharedBuffer(mBuffer, offset));

        mAttributeList.push_back(pAttribute);
        offset += pAttr->Length;
    }
}

CFileRecord::~CFileRecord()
{                  
    for (unsigned int i=0; i < mAttributeList.size(); i++)
        delete mAttributeList[i];
}

const CAttribute* CFileRecord::getAttribute(ATTRIBUTE_TYPE AttributeType) const
{ 
    ATTRIBUTE_TYPE at;
    for (unsigned int i = 0; i < getAttributesCount(); i++)
    {
        CAttribute* pAttr = const_cast<CAttribute*>( getAttribute(i) );
        at = PATTRIBUTE(*pAttr)->AttributeType;
        if (at == AttributeType)
            return pAttr;
    }
    return NULL;
}

/******************************************************************************/
/*                                   CMFT                                     */
/******************************************************************************/

CMFT::CMFT(CNTFSVolume* pVolume): 
    CFileRecord(), mpVolume(pVolume)
{
    SxSharedBuffer Buff(mpVolume->getFileRecordLength());
    mBuffer = Buff;
    mpVolume->read(mBuffer, mpVolume->getFileRecordLength(), mpVolume->getMFTOffset());
    USAFixupFileRecord(mBuffer);
    LoadAttributes();
}

CFileRecord CMFT::getFileRecord(UINT64 Index)
{
    UINT64 VCN = Index / (mpVolume->getClusterLength() / mpVolume->getFileRecordLength());
    INT32  subIndex = Index % (mpVolume->getClusterLength() / mpVolume->getFileRecordLength());

    const CNonResidentAttribute* nrAttr = dynamic_cast<CNonResidentAttribute*>(getAttribute(AttributeData));
    if (!nrAttr)
        throw string("No data attribute");
    UINT64 LCN = nrAttr->getLCN(VCN);
    UINT64 readOffset = LCN*mpVolume->getClusterLength() + subIndex*mpVolume->getFileRecordLength();

    SxSharedBuffer Buff(mpVolume->getFileRecordLength());
    mpVolume->read(Buff, mpVolume->getFileRecordLength(), readOffset);
    USAFixupFileRecord(Buff);
    
    return CFileRecord(Buff);
}


/******************************************************************************/
/*                                  CAttribure                                */
/******************************************************************************/

SxSharedBuffer getAttributeValue(const CResidentAttribute& Attribute)
{
    PCRESIDENT_ATTRIBUTE attr = Attribute.mBuffer;
    SxSharedBuffer mem(attr->ValueLength);

    memcpy(mem, IncPTR(attr, attr->ValueOffset), attr->ValueLength);
    return mem;
}

SxSharedBuffer getAttributeValue(const CNonResidentAttribute& Attribute, const CNTFSVolume& Volume) 
{
    PCNONRESIDENT_ATTRIBUTE attr = Attribute.mBuffer;
    SxSharedBuffer mem(UINT32(attr->DataSize));

    for (UINT64 vcn = attr->LowVcn; vcn <= attr->HighVcn; vcn++ )
    {
        SxSharedBuffer sector = Volume.readCluster(Attribute.getLCN(vcn));

        UINT32 cbLeft = UINT32(attr->DataSize - vcn*Volume.getClusterLength());
        UINT32 cbSize = (cbLeft > Volume.getClusterLength())?Volume.getClusterLength():cbLeft;

        memcpy(IncPTR((void*)mem, (UINT32)vcn*Volume.getClusterLength()), sector, cbSize);
    }
    
    return mem;
}

/******************************************************************************/
/*                                  CNTFSRun                                  */
/******************************************************************************/

CNTFSRun::CNTFSRun(const PCNONRESIDENT_ATTRIBUTE pNonResidentAttribute)
{
    RUNITEM runItem;
    PUINT8 pRun = (PUINT8)IncPTR(pNonResidentAttribute, pNonResidentAttribute->RunArrayOffset);
    LowVCN = pNonResidentAttribute->LowVcn;
    HighVCN = pNonResidentAttribute->HighVcn;
    runItem.LowVCN = pNonResidentAttribute->LowVcn;

    while (*pRun)
    {
        UINT8 cbRunClusersCount = *pRun & 0x0F;
        UINT8 cbRunStartLCN = (*pRun >> 4) & 0x0F;
        UINT8 runLength =  cbRunClusersCount + cbRunStartLCN + 1;

        runItem.ClustersCount = BytesToInt64((pRun+1), cbRunClusersCount);
        runItem.LowLCN = BytesToInt64((pRun+cbRunClusersCount+1), cbRunStartLCN);
        runItem.HighVCN = runItem.LowVCN + runItem.ClustersCount; 
        runItem.HighLCN = runItem.LowLCN + runItem.ClustersCount; 
        mRunList.push_back(runItem);
        runItem.LowVCN = runItem.HighVCN + 1;
        pRun = IncPTR(pRun, runLength);
    }
}

UINT64 CNTFSRun::LCN(UINT64 VCN) const
{
    for (unsigned int i = 0; i < mRunList.size(); i++)
    {
        if ( (VCN >= mRunList[i].LowVCN) && (VCN <= mRunList[i].HighVCN) )
            return mRunList[i].LowLCN + (VCN - mRunList[i].LowVCN);
    }
    return UINT64(-1);
}

}  //  namespace ntfs
